"""A simple xlsx file reader, supporting a very limited subset of the format, designed to
be as similar to xld module API as possible

"""
USE_EL_TREE = 1 # NB: Writing requires this
if USE_EL_TREE:
    import xml.etree.cElementTree as ET
    #from ET import parse
    parse = ET.parse
    USE_SHARED_STRING = 1
    import xml
    nsS = '{http://schemas.openxmlformats.org/spreadsheetml/2006/main}'
    nsRP = '{http://schemas.openxmlformats.org/package/2006/relationships}'
    nsR = '{http://schemas.openxmlformats.org/officeDocument/2006/relationships}'
else:
    from xml.dom.minidom import parse
import zipfile
class xlsx:
    def __init__(self, fpath, write = False):
        """Open the zip file, and set up lazy caches for the xml files and specific sheets"""
        self.writeable = write
        if write:
            self.zipFilePath = fpath
            self.zf = zipfile.ZipFile(fpath,'r')
            self.zipFiles = {} # for our eventual copy
        else:
            self.zf = zipfile.ZipFile(fpath,'r')
        self.xmlFiles = {}
        self.sheets = {}
        self.sharedStrings = {}
        self.sharedStringsDirty = False
    def xml(self, name):
        """Extra a parsed DOM of the given XML file in the archive"""
        if self.xmlFiles.has_key(name):
            return self.xmlFiles[name]
        try:
            f = self.zf.open(name)
        except:
            return None
        dom = parse(f)
        if USE_EL_TREE:
            dom = dom.getroot()
        self.xmlFiles[name] = dom
        return dom
    def sharedStringElement(self):
        """Get the xlsx shared string file dom"""
        return self.xml('xl/sharedStrings.xml')
    def sharedStringWithIndex(self, index):
        """Find a given shared string by index - we end up tossing any formmating"""
        if self.sharedStrings.has_key(index):
            return self.sharedStrings[index]
        if USE_EL_TREE:
            self.sharedStrings[index] = self.sharedStringElement().findall(nsS + 'si')[index][0].text
        else:
            self.sharedStrings[index] = self.sharedStringElement().getElementsByTagName('si')[index].firstChild.firstChild.toxml()
        return self.sharedStrings[index]
    def addSharedString(self, string):
        if USE_EL_TREE:
            retval = 0
            sst = self.sharedStringElement()
            for si in sst.findall(nsS + 'si'):
                if si[0].text == string:
                    return retval # already there
                retval += 1
            #print "Adding shared string #" + str(retval), string
            self.sharedStrings[retval] = string
            self.sharedStringsDirty = True
            sst.set("count", str(int(sst.get("count",0))+1))
            sst.set("uniqueCount", str(int(sst.get("uniqueCount",0))+1))
            inlineString = ET.SubElement(sst, nsS + "si")
            ET.SubElement(inlineString, nsS + "t").text = string
            return retval
        return 0
            
    def sheet_by_name(self, name):
        """Find a specific sheet in the xlsx file by name"""
        if not self.sheets:
            workbook = self.xml("xl/workbook.xml")
            workbookRels = self.xml("xl/_rels/workbook.xml.rels")
            relMap = {}
            if USE_EL_TREE:
                for rel in workbookRels.findall(nsRP + 'Relationship'):
                    relMap[rel.get('Id')] = rel.get('Target')
                for sheets in workbook.findall(nsS + 'sheets'):
                    for sheet in sheets.findall(nsS + 'sheet'):
                        #print sheet.toxml()
                        #                    print sheet.attributes._attrs
                        childName = relMap[sheet.get(nsR + 'id')]
                        sheetName = sheet.get('name')
                        self.sheets[sheetName] = { 'path':childName }
            else:
                for rel in workbookRels.getElementsByTagName('Relationship'):
                    relMap[rel.attributes['Id'].value] = rel.attributes['Target'].value
                for sheets in workbook.getElementsByTagName('sheets'):
                    for sheet in sheets.getElementsByTagName('sheet'):
                        #print sheet.toxml()
                        #                    print sheet.attributes._attrs
                        childName = relMap[sheet.attributes['r:id'].value]
                        sheetName = sheet.attributes['name'].value
                        self.sheets[sheetName] = { 'path':childName }
        value = self.sheets[name]
        if not value.has_key('sheet'):
            sheetXML = self.xml('xl/' + value['path'])
            # print "Loading sheet", name, ':', value['path']
            sheet = xlsxSheet(sheetXML, self, value['path'], name)
            value['sheet'] = sheet
            # try to load rels
            relsPathParts = value['path'].split('/')
            relsPathParts.insert(-1,'_rels')
            relsPathParts[-1] = relsPathParts[-1] + ".rels"
            relXML = self.xml('xl/' + '/'.join(relsPathParts))
            value['rels' ] = relXML
        return value['sheet']
    def relsForSheet(self, sheet):
        for value in self.sheets.values():
            if value.get('sheet') == sheet:
                return value.get('rels')
        return None
    def saveSheet(self, sheet):
        # remove the calc chain
        # This is ugly because it is global, but does what we need
        ET.register_namespace("",nsS[1:-1])
        ET.register_namespace("r",nsR[1:-1])
        ET.register_namespace("xdr","http://schemas.openxmlformats.org/drawingml/2006/spreadsheetDrawing")
        ET.register_namespace("x14","http://schemas.microsoft.com/office/spreadsheetml/2009/9/main")
        ET.register_namespace("mc","http://schemas.openxmlformats.org/markup-compatibility/2006")
        ET.register_namespace("x14ac","http://schemas.microsoft.com/office/spreadsheetml/2009/9/ac")
        # these two appear within elements, and not global to the document - do the best we can
        ET.register_namespace("xm","http://schemas.microsoft.com/office/excel/2006/main")
        ET.register_namespace("mx","http://schemas.microsoft.com/office/mac/excel/2008/main")
        xmlstr = ET.tostring(sheet.xml, encoding="UTF-8")
        self.zipFiles['xl/' + sheet.path] = xmlstr
        sheet.dirty = False
    def close(self):
        # by default, overwrite where we came from
        if self.writeable:
            self.write(self.zipFilePath)
    def write(self, destPath):
        if not self.writeable:
            return
        # we need to unpack the entire file and rezip without duplicate entries for Excel
        # (this also makes it easier to edit an existing file)
        if not self.zipFiles:
            self.zipFiles = {}
            fileList = self.zf.namelist()
            for fileName in fileList:
                self.zipFiles[fileName] = self.zf.open(fileName).read()
        for sheet in self.sheets.values():
            if sheet.has_key('sheet'):
                sheet['sheet'].save()
        if self.sharedStringsDirty:
            xmlstr = ET.tostring(self.sharedStringElement(), encoding="UTF-8")
            self.zipFiles['xl/sharedStrings.xml'] = xmlstr
        # now make a new file here
        # self.zf.close() # we don't need to close it, the file reference remains valid
        newZipFile = zipfile.ZipFile(destPath,"w")
        for fileName in fileList:
            newZipFile.writestr(fileName, self.zipFiles[fileName])
        newZipFile.close()
class xlsxSheet:
    """Represents a specific sheet in the workbook"""
    def __init__(self, xml, document, path, name):
        self.xml = xml
        self.name = name
        self.book = document
        self.rows = []
        self.path = path
        self.dirty = False
    def allComments(self):
        """Get all the comments for this sheet"""
        rels = self.book.relsForSheet(self)
        if not rels:
            return None
        if USE_EL_TREE:
            import os
            retval = {}
            for rel in rels.findall(nsRP + 'Relationship'):
                if rel.get('Type') == "http://schemas.openxmlformats.org/officeDocument/2006/relationships/comments":
                    baseDir = os.path.split(self.path)[0]
                    commentPath = os.path.normpath(os.path.join(baseDir,rel.get('Target')))
                    commentsXML = self.book.xml('xl/' + commentPath)
                    commentList = commentsXML.find(nsS+'commentList')
                    for comment in commentList.findall(nsS+'comment'):
                        commentText = ""
                        for t in comment.findall(nsS+'text'):
                            for i in t.itertext():
                                commentText += i
                        retval[comment.get('ref')] = commentText
            return retval
    def row(self, rowIndex):
        """Get a specific row in the sheet"""
        if not self.rows:
            self.sharedFormulas = []
            self.rowsXML = {}
            if USE_EL_TREE:
                for el in self.xml.find(nsS + 'sheetData').findall(nsS + 'row'):
                    rnum = int(el.get('r'))
                    rlen = len(el.get('r'))
                    self.rowsXML[rnum] = el
                    #print "found row", el, "num", rnum, "len", rlen
                    while len(self.rows) <= rnum:
                        self.rows.append({})
                    cols = self.rows[rnum]
                    for elc in el.findall(nsS + 'c'):
                        rname = elc.get('r')
                        cols[rname[:-rlen]] = elc
            else:
                for el in self.xml.getElementsByTagName('sheetData')[0].getElementsByTagName('row'):
                    rnum = int(el.attributes['r'].value)
                    rlen = len(el.attributes['r'].value)
                    while len(self.rows) <= rnum:
                        self.rows.append({})
                    cols = self.rows[rnum]
                    for elc in el.getElementsByTagName('c'):
                        rname = elc.attributes['r'].value
                        cols[rname[:-rlen]] = elc
# check for shared formulas
        return self.rows[rowIndex]
    def colName(self, colIndex):
        """Convert a colIndex (0..25, etc...) as A..Z, AA..ZZ"""
        if colIndex < 26:
            return chr(ord('A') + colIndex)
        return chr(ord('A') + colIndex / 26 - 1) + chr(ord('A') + colIndex % 26)
    def cell(self, rowIndex, colIndex):
        """Get the data of a specific cell by row and column (zero based)"""
        row = self.row(rowIndex+1)
        colName = self.colName(colIndex)
        if not row.has_key(colName):
            return empty_cell
        cell = row[colName]
        if USE_EL_TREE:
            cellType = cell.get('t')
            if cellType == 'inlineStr':
                str = ""
                for isXML in cell.findall(nsS + "is"):
                    for tXML in isXML.findall(nsS + "t"):
                        str += tXML.text
                return xlsxValue(XL_CELL_TEXT, str)
        if USE_EL_TREE:
            if cell.get('si'):
                print "Cell",colName,rowIndex+1,"has si:",cell.text
            v = cell.findall(nsS + 'v')
        else:
            if cell.attributes.has_key('si'):
                print "Cell",colName,rowIndex+1,"has si:",cell.toxml()
            v = cell.getElementsByTagName('v')
        if not v:       
            #            print "No V for:",cell.toxml()
            return empty_cell
        if USE_EL_TREE:
            v = v[0].text
            if cell.get('t'):
                if cellType == 's':
                    v = int(v)
                    return xlsxValue(XL_CELL_TEXT,self.book.sharedStringWithIndex(v))
                if cellType == 'str': # a formula, but grab latest value
                    return xlsxValue(XL_CELL_TEXT, v)
    #            print "cell",rowIndex, colIndex, cellType, v
        else:
            v = v[0].firstChild.toxml()
            if cell.attributes.has_key('t'):
                cellType = cell.attributes['t'].value
                if cellType == 's':
                    v = int(v)
                    return xlsxValue(XL_CELL_TEXT,self.book.sharedStringWithIndex(v))
                if cellType == 'str':
                    return xlsxValue(XL_CELL_TEXT, v)
    #            print "cell",rowIndex, colIndex, cellType, v
        return xlsxValue(XL_CELL_NUMBER,float(v))
    def save(self):
        if self.dirty:
            print "Saving",self.name,"as", self.path
            self.book.saveSheet(self)
    def setCell(self, rowIndex, colIndex, value):
        self.dirty = True
        clearing = value == None or value == ""
        # a very limited ability to set cells
        row = self.row(rowIndex+1)
        rowXML = self.rowsXML[rowIndex+1]
        if not rowXML:
            raise "Missing row xml for writing"
        colName = self.colName(colIndex)
        if not row.has_key(colName):
            if clearing:
                return # we are deleting the value, no cell -> nothing to do
            # need to add col
            #print "Adding col", colName, rowIndex + 1, row
            colXML = ET.Element(nsS + "c")
            colXML.set("r", "%s%d" % (colName, rowIndex+1))
            rowXML.append(colXML)
            row[colName] = colXML
        else:
            colXML = row[colName]
        v = colXML.findall(nsS + 'v')
        #colXML.attrib.pop("s", None) # remove the "s" attribute???
        if clearing:
            # if there is a "v", clear it
            if v:
                v = v[0]
                colXML.remove(v)
                colXML.pop("t", None) # is this correct?
        elif type(value) == type("") or type(value) == type(u""):
            if colXML.get("t") == "str": # this is a formula with string value (such as jammer number)
                # so just fill in the str value here
                if not v:
                    v = ET.Element(nsS +  "v")
                    colXML.append(v)
                else:
                    v = v[0]
                v.text = str(value)
                return
            if USE_SHARED_STRING:
                if not v:
                    v = ET.Element(nsS +  "v")
                    colXML.append(v)
                else:
                    v = v[0]
                colXML.set("t","s")
                v.text = str(self.book.addSharedString(value))
            else:
                # clear v
                if v:
                    v = v[0]
                    colXML.remove(v)
                colXML.set("t", "inlineStr")
                # see if this fixes the inlineStr problem with lineup
                if value == "$":
                    value = "D"
                elif value == "/":
                    value = "F"
                inlineString = ET.SubElement(colXML, nsS + "is")
                ET.SubElement(inlineString, nsS + "t").text = str(value)
        elif type(value) == type(1) or type(value) == type(1.0):
            if not v:
                v = ET.Element(nsS +  "v")
                colXML.append(v)
            else:
                v = v[0]
            v.text = str(value)
            colXML.set("t", "n")


class xlsxValue:
    """Represents the contents of a given cell in a sheet, as a type and value"""
    def __init__(self,type,value):
        self.type = type
        self.value = value
    def __repr__(self):
        if self.type == XL_CELL_EMPTY:
            return "<EMPTY>"
        if self.type == XL_CELL_TEXT:
            return '"' + self.value + '"'
        if self.type == XL_CELL_NUMBER:
            return "%g" % (self.value,)
        if self.type == XL_CELL_ERROR:
            return "<ERROR " + self.value + ">"
        return "<???>"
XL_CELL_EMPTY = 0
XL_CELL_TEXT = 1
XL_CELL_NUMBER = 2
XL_CELL_DATE = 3
XL_CELL_BOOLEAN = 4
XL_CELL_ERROR = 5
XL_CELL_BLANK = 6
empty_cell = xlsxValue(XL_CELL_EMPTY,u'')

def open_workbook(f, write = False):
    return xlsx(f, write)
if __name__ == '__main__':
#    x = xlsx('/Users/gandreas/Downloads/STATS-110513_FairbanksRollerGirls_vs_RageCityRollergirls.xlsx')
    x = xlsx('/Users/gandreas/Documents/STATS-2014-04-27-NorthStarRollerGirls_vs_OldCapitolCityRollerGirls.xlsx')
    #print x.sharedStringElement().toxml()
    #print x.sharedStringWithIndex(5).toxml()
    IBRF = x.sheet_by_name("IGRF")
    #print IBRF.row(3)['B'].toxml()

    print IBRF.cell(2,1)
    print IBRF.cell(4,1)
    print IBRF.cell(10,1)
    print IBRF.allComments()
    Lineups = x.sheet_by_name('Lineups')
    print Lineups.cell(0,0)
    
